# coding:utf-8
#!/usr/bin/env python
from PoboAPI import *
import datetime
import statsmodels.formula.api as sm
from statsmodels.tsa import stattools
import numpy as np
import pandas as pd
from copy import *

# Python3Preview
#开始时间，用于初始化一些参数        
def OnStart(context) :
    print("I\'m starting...")
    #获取ag当前可交易合约，并按月份进行排序
    vac=sorted(GetFuturesContracts('SHFE','rb'))
    #获取当前主力合约
    MainContract=GetMainContract('SHFE', 'rb',20)
    #获取当前主力合约下一月合约
    for i in range(len(vac)):
        if vac[i]==MainContract:
            s=i+1
            break
    g.code=[MainContract,vac[s]]
    #订阅g.code[0]
    SubscribeQuote(g.code[0])
    #登录交易账号
    if "回测期货" in context.accounts :
        print("登录交易账号 回测期货")
        if context.accounts["回测期货"].Login() :
            context.myacc = context.accounts["回测期货"]

#实时行情事件，当有新行情出现时调用该事件                        
def OnQuote(context,code) :
    #print "调用到OnQuote事件"
    vac=sorted(GetFuturesContracts('SHFE','rb'))
    #获取当前主力合约
    MainContract=GetMainContract('SHFE', 'rb',20)
    #获取当前主力合约下一月合约
    price_type=PriceType(PbPriceType.LimitFOK,limit_price_type=2)
    for i in range(len(vac)):
        if vac[i]==MainContract:
            s=i+1
            break
    if g.code[0]!=MainContract:
        #进行平仓
        for c in g.code:
            option = PBObj()
            option.contract = c
            pos = context.myacc.GetPositions(option)
            for p in pos:
                if p.bstype.BuySellFlag=="0":
                    QuickInsertOrder(context.myacc,c,'sell','close',price_type,p.volume)
                if p.bstype.BuySellFlag=="1":
                    QuickInsertOrder(context.myacc,c,'buy','close',price_type,p.volume)
        #取消订阅原主力合约
        UnsubscribeQuote(g.code[0])
        g.code_old=g.code
        g.code=[MainContract,vac[s]]
        print("当前主力合约由"+str(g.code_old[0])+"更换为"+str(g.code[1]))
        print("平仓掉原主力合约")
        #订阅当前主力合约信息
        SubscribeQuote(g.code[0])
    print("当前套利组合为"+str(g.code))
    #获取当前日期
    current_date=GetCurrentTradingDate('SHFE')
    #获取两个品种的近60个交易日的收盘价构造收盘价序列
    close_01= GetHisDataByFieldAsDF(g.code[0],'close', BarType.Day, start_date=None, end_date=current_date, count=60,weight_type=0)
    close_02= GetHisDataByFieldAsDF(g.code[1],'close', BarType.Day, start_date=None, end_date=current_date, count=60,weight_type=0)
    #对价格序列取对数
    x=np.log(close_01)
    y=np.log(close_02)
    #做回归，取残差序列
    x=pd.Series(x)
    y=pd.Series(y)
    df = pd.concat([x,y],axis = 1)
    df.columns = ['x','y']
    model = sm.ols(formula = 'y~x',data = df).fit()
    resid=model.resid
    #对残差进行单位根检验，求p值
    p=stattools.adfuller(resid)[1]
    #设定残差序列的2倍标准差为开仓阈值
    open_ththreshold_positive =2*resid.std()
    #设定残差序列的0.01倍标准差为平仓阈值
    close_ththreshold_positive =0.05*resid.std()
    close_ththreshold_negative =-0.05*resid.std()
    resid=np.array(resid)
    #获取资产的当前涨停价，成交量，最新价
    dyndata1 = GetQuote(g.code[0])
    dyndata2 = GetQuote(g.code[1])
    riselimit1=dyndata1.riselimit
    riselimit2=dyndata2.riselimit
    value = 0.8
    context.myacc.SetRiskDegreeMaxThreshold(value)
    isOver=context.myacc.IsRiskDegreeOverMax()
    if isOver==False:
        #若通过协整检验
        if p<0.05:
            value = PBObj()
#             value.Rate = 0.12
#             context.myacc.SetMargin(g.code[0], value)
#             context.myacc.SetMargin(g.code[1], value)
            #下单
            if resid[-1]>= open_ththreshold_positive:
                v1=context.myacc.GetValidVolume(g.code[0],BSType.BuyOpen,riselimit1)
                v2=context.myacc.GetValidVolume(g.code[1],BSType.SellOpen,riselimit2)
                v=int(min(v1,v2))
                QuickInsertOrder(context.myacc,g.code[0],'buy','open',price_type,v//2)
                QuickInsertOrder(context.myacc,g.code[1],'sell','open',price_type,v//2)
                print("买入"+str(g.code[0])+","+"卖出"+str(g.code[1]))
            #平仓
            if resid[-1]<close_ththreshold_positive:
                for c in g.code:
                    option = PBObj()
                    option.contract = c
                    pos = context.myacc.GetPositions(option)
                    dyndata=GetQuote(c)
                    now=dyndata.now
                    for p in pos:
                        if p.bstype.BuySellFlag=="0":
                            QuickInsertOrder(context.myacc,c,'sell','close',price_type,p.volume)
                        print("卖平"+str(c))
                        if p.bstype.BuySellFlag=="1":
                            QuickInsertOrder(context.myacc,c,'buy','close',price_type,p.volume)
                        print("买平"+str(c))

